using BenchmarkTools
using JuMP
using SparseVariables

const SV = SparseVariables

# Demo custom types
abstract type Node end
struct Source <: Node
    id::Symbol
end
Source(n::Int) = Source(Symbol(string(n)))
struct Sink <: Node
    id::Symbol
end
Sink(n::Int) = Sink(Symbol(string(n)))

# Translation to type-stable index
@inline SV.translate(n::Node) = n.id
@inline SV.translate(::Type{<:Node}) = Symbol

"""
    Indexing on custom types with `IndexedVarArray`
"""
function test(N)
    m = Model()
    @variable(
        m,
        x[i = Sink.(1:N), y = Source.(1:N)],
        container = SV.IndexedVarArray
    )
    for i in Sink.(1:N), j in Source.(1:N)
        insertvar!(x, i, j)
    end
end

"""
    Indexing on `Int` manually for performance with `IndexedVarArray`
"""
function test_ix(N)
    m = Model()
    @variable(m, x[i = 1:N, j = 1:N], container = SV.IndexedVarArray)
    for i = 1:N, j = 1:N
        insertvar!(x, i, j)
    end
end

"""
    Automatic translation to `Int` by `TranslateVarArray`
"""
function test_t(N)
    m = Model()
    @variable(
        m,
        x[i = Sink.(1:N), j = Source.(1:N)],
        container = SV.TranslateVarArray
    )
    for i in Sink.(1:N), j in Source.(1:N)
        insertvar!(x, i, j)
    end
end

"""
    Manual indexing by `Int` with `TranslateVarArray` for comparison
"""
function test_i(N)
    m = Model()
    @variable(m, x[i = 1:N, j = 1:N], container = SV.TranslateVarArray)
    for i = 1:N, j = 1:N
        insertvar!(x, i, j)
    end
end


N = 100
@info "Custom type w/IndexedVarArray"
@btime test(N)

@info "Manual Int w/IndexedVarArray"
@btime test_ix(N)

@info "Automatic w/TranslateVarArray"
@btime test_t(N)

@info "Manual Int w/TranslateVarArray"
@btime test_i(N)